#!/usr/bin/python

__author__ = "Jimmy Saw"
__copyright__ = "Copyright 2011, extracts intergenic sequences"
__credits__ = "Jimmy Saw"
__email__ = "jimmysaw@gmail.com"
__updated__ = "01/05/2012"

"""
TO DO: 
Usage: python auto_anno.py annofile.txt repeat_regions.list seqfile.fasta prefix minlen
annofile.txt (tab-delimited file with 6 fields)
repeat_region.list (tab-delimited file with repeat coordinates)
seqfile.fasta (genome fasta file)
prefix (name for output prefix)
minlen (minimum length of IGS to extract)

"""

import sys
from Bio import SeqIO
from Bio.SeqRecord import SeqRecord

#Formatting for 4 significant figures (for decimals)
def fmt(f):
    st = '{0:.4}'.format(f)
    return st

#Check start sites begins
annofile = sys.argv[1]
repeats = sys.argv[2]
seqfile = sys.argv[3]
prefix = sys.argv[4]
minlen = int(sys.argv[5])
af = open(annofile, "rU")
rf = open(repeats, "rU")
sf = open(seqfile, "rU")
rec = SeqIO.read(sf, "fasta")
lines = af.readlines()
rlines = rf.readlines()
num = len(lines)

repeat_start = []
repeat_stop = []

#list of repeat features to avoid extracting
for line in rlines:
    l = line.split('\t')
    repeat_start.append(int(l[0]))
    repeat_stop.append(int(l[1].rstrip()))

i = 0

while i < num:
    if i == 0:
        curr_line = lines[i].split('\t')
        curr_id = curr_line[0]
        curr_start = int(curr_line[3])
        substop = curr_start - 1
        for x, y in enumerate(repeat_start):
            if repeat_start[x] < curr_start < repeat_stop[x] and repeat_start[x] < substop < repeat_stop[x]:
                curr_start = 0
                substop = 0
            elif repeat_start[x] < curr_start < repeat_stop[x] or repeat_start[x] < substop < repeat_stop[x]:
                curr_start = 0
                substop = 0
            else:
                pass
        if substop > curr_start:
            igs_id = prefix + "_IGS_" + str(i).zfill(4) + "_" + str(substart) + "-" + str(substop)
            subseq = rec.seq[curr_start:substop]
            if len(subseq) > minlen:
                tmprec = SeqRecord(subseq, id=igs_id, description="")
                print tmprec.format("fasta").strip('\n')
                #print ">" + igs_id
                #print subseq                
    if i == num - 1:
        curr_line = lines[i].split('\t')
        curr_stop = int(curr_line[4])
        substart = curr_stop + 1
        substop = len(rec.seq)
        for x, y in enumerate(repeat_start):
            if repeat_start[x] < substart < repeat_stop[x] and repeat_start[x] < substop < repeat_stop[x]:
                substart = 0
                substop = 0
            elif repeat_start[x] < substart < repeat_stop[x] or repeat_start[x] < substop < repeat_stop[x]:
                substart = 0
                substop = 0
            else:
                pass
        if substop > substart:
            igs_id = prefix + "_IGS_" + str(i).zfill(4) + "_" + str(substart) + "-" + str(substop)
            subseq = rec.seq[substart:substop]
            if len(subseq) > minlen:
                tmprec = SeqRecord(subseq, id=igs_id, description="")
                print tmprec.format("fasta").strip('\n')
                #print ">" + igs_id
                #print subseq
    if i != num - 1 and i != 0:
        curr_line = lines[i].split('\t')
        curr_id = curr_line[0]
        curr_locus_tag = curr_line[1]
        curr_feat_type = curr_line[2]
        curr_start = int(curr_line[3])
        curr_stop = int(curr_line[4])
        curr_frame = curr_line[5]

        next_line = lines[i + 1].split('\t')
        next_id = next_line[0]
        next_locus_tag = next_line[1]
        next_feat_type = next_line[2]
        next_start = int(next_line[3])
        next_stop = int(next_line[4])
        next_frame = next_line[5]

        prev_line = lines[i - 1].split('\t')
        prev_id = prev_line[0]
        prev_locus_tag = prev_line[1]
        prev_feat_type = prev_line[2]
        prev_start = int(prev_line[3])
        prev_stop = int(prev_line[4])
        prev_frame = prev_line[5]

        #extracts IGS only upstream of features that are CDS thus avoiding having to deal with rRNA or other genes
        if curr_feat_type == "CDS":
            substart = prev_stop + 1
            substop = curr_start - 1
            for x, y in enumerate(repeat_start):
                if repeat_start[x] < substart < repeat_stop[x] and repeat_start[x] < substop < repeat_stop[x]:
                    substart = 0
                    substop = 0
                elif repeat_start[x] < substart < repeat_stop[x] or repeat_start[x] < substop < repeat_stop[x]:
                    substart = 0
                    substop = 0
                else:
                    pass
            if substop > substart:
                igs_id = prefix + "_IGS_" + str(i).zfill(4) + "_" + str(substart) + "-" + str(substop)
                subseq = rec.seq[substart:substop]
                if len(subseq) > minlen:
                    tmprec = SeqRecord(subseq, id=igs_id, description="")
                    print tmprec.format("fasta").strip('\n')
                    #print ">" + igs_id
                    #print subseq
    i += 1



af.close()
sf.close()
